import numpy as np
import pandas as pd
from python_ai.common.xcommon import sep
pd.set_option('display.max_columns', None)

from sklearn.feature_extraction.text import CountVectorizer

df = pd.read_csv('bayes_xinxi.txt')

tf_model = CountVectorizer(token_pattern=r'[a-zA-Z]+')
X = df.iloc[:, 1]
Y = df.iloc[:, 0]
sep('X')
print(X)

# X = tf_model.fit_transform(X).A  # ATTENTION .A is not necessary! Sparse matrix is more efficient.
X = tf_model.fit_transform(X)
sep('feature names')
print(tf_model.get_feature_names())
print('X after fit trans')
print(X)

X_ = ['Chinese Chinese Chinese Tokyo Japan']
# X_ = tf_model.transform(X_).A  # ATTENTION .A is not necessary! Sparse matrix is more efficient.
X_ = tf_model.transform(X_)  # ATTENTION Here, just transform!
sep('X sample')
print(X_)

from sklearn.naive_bayes import MultinomialNB
model = MultinomialNB()
model.fit(X, Y)
print(model.predict(X_))
